import time
from machine import Pin, PWM

class Gait:
    def __init__(self, leg_pins):
        """
        初始化六足机器人步态控制器
        leg_pins: 包含6个腿的PWM引脚信息的列表，每个腿包含3个舵机的引脚
        """
        self.legs = []
        self.leg_pins = leg_pins
        self.gait_type = "tripod"  # 默认使用三角步态
        self.direction = "forward"  # 默认前进方向
        self.speed = 50  # 默认速度
        
        # 初始化所有舵机
        self.init_servos()
        
        # 定义腿部组（三角步态使用）
        self.tripod_group1 = [0, 2, 4]  # 左前、右中、左后
        self.tripod_group2 = [1, 3, 5]  # 右前、左中、右后
        
        # 定义步态参数
        self.step_height = 30  # 抬腿高度
        self.step_length = 20  # 步长
        self.center_positions = [90, 90, 90] * 6  # 所有舵机的中心位置
        
    def init_servos(self):
        """初始化所有舵机PWM输出"""
        for leg_idx, pins in enumerate(self.leg_pins):
            leg_servos = []
            for pin in pins:
                pwm = PWM(Pin(pin))
                pwm.freq(50)  # 舵机频率通常为50Hz
                leg_servos.append(pwm)
            self.legs.append(leg_servos)
    
    def set_servo_position(self, leg_index, servo_index, angle):
        """设置指定舵机的角度"""
        # 确保角度在有效范围内
        angle = max(0, min(180, angle))
        # 将角度转换为PWM占空比 (2000-7800对应0-180度)
        duty = int(2000 + (angle / 180) * 5800)
        self.legs[leg_index][servo_index].duty_u16(duty)
    
    def set_leg_position(self, leg_index, angles):
        """设置整条腿的三个舵机角度"""
        for i, angle in enumerate(angles):
            self.set_servo_position(leg_index, i, angle)
    
    def tripod_gait(self, phase):
        """三角步态实现"""
        # 确定当前处于哪一组抬腿阶段
        if phase < 0.5:
            active_legs = self.tripod_group1
            support_legs = self.tripod_group2
        else:
            active_legs = self.tripod_group2
            support_legs = self.tripod_group1
        
        # 归一化相位到0-1之间
        norm_phase = phase * 2 % 1
        
        # 根据方向调整步态
        if self.direction == "forward":
            step_offset = self.step_length * (0.5 - norm_phase)
        elif self.direction == "backward":
            step_offset = self.step_length * (norm_phase - 0.5)
        elif self.direction == "left":
            # 左转实现
            step_offset = [
                self.step_length * (0.5 - norm_phase),  # 右前腿前进
                self.step_length * (norm_phase - 0.5),  # 左前腿后退
                self.step_length * (0.5 - norm_phase),  # 右中腿前进
                self.step_length * (norm_phase - 0.5),  # 左中腿后退
                self.step_length * (0.5 - norm_phase),  # 右后腿前进
                self.step_length * (norm_phase - 0.5)   # 左后腿后退
            ]
        elif self.direction == "right":
            # 右转实现
            step_offset = [
                self.step_length * (norm_phase - 0.5),  # 右前腿后退
                self.step_length * (0.5 - norm_phase),  # 左前腿前进
                self.step_length * (norm_phase - 0.5),  # 右中腿后退
                self.step_length * (0.5 - norm_phase),  # 左中腿前进
                self.step_length * (norm_phase - 0.5),  # 右后腿后退
                self.step_length * (0.5 - norm_phase)   # 左后腿前进
            ]
        
        # 为抬腿组创建抬腿动作
        lift_height = self.step_height * (0.5 - abs(norm_phase - 0.5)) * 2
        
        # 更新所有腿的位置
        for leg in range(6):
            # 计算基础位置
            base_positions = self.center_positions[leg*3:(leg+1)*3].copy()
            
            # 应用步态偏移
            if self.direction in ["left", "right"]:
                # 左右转使用数组形式的step_offset
                base_positions[0] += step_offset[leg]
            else:
                # 前进后退使用统一的step_offset
                base_positions[0] += step_offset
            
            # 抬腿动作
            if leg in active_legs:
                base_positions[1] += lift_height
            
            # 设置腿的位置
            self.set_leg_position(leg, base_positions)
    
    def wave_gait(self, phase):
        """波浪形步态实现"""
        # 波浪步态中，每条腿的相位差为1/6周期
        for leg in range(6):
            leg_phase = (phase + leg/6) % 1
            
            # 根据方向调整步态
            if self.direction == "forward":
                step_offset = self.step_length * (0.5 - leg_phase)
            elif self.direction == "backward":
                step_offset = self.step_length * (leg_phase - 0.5)
            elif self.direction == "left":
                # 左转实现
                if leg % 2 == 0:  # 偶数索引腿(右)前进
                    step_offset = self.step_length * (0.5 - leg_phase)
                else:  # 奇数索引腿(左)后退
                    step_offset = self.step_length * (leg_phase - 0.5)
            elif self.direction == "right":
                # 右转实现
                if leg % 2 == 0:  # 偶数索引腿(右)后退
                    step_offset = self.step_length * (leg_phase - 0.5)
                else:  # 奇数索引腿(左)前进
                    step_offset = self.step_length * (0.5 - leg_phase)
            
            # 创建抬腿动作
            lift_height = self.step_height * (0.5 - abs(leg_phase - 0.5)) * 2
            
            # 获取基础位置
            base_positions = self.center_positions[leg*3:(leg+1)*3].copy()
            
            # 应用步态偏移
            base_positions[0] += step_offset
            base_positions[1] += lift_height
            
            # 设置腿的位置
            self.set_leg_position(leg, base_positions)
    
    def set_gait_type(self, gait_type):
        """设置步态类型：'tripod' 或 'wave'"""
        if gait_type in ["tripod", "wave"]:
            self.gait_type = gait_type
        else:
            print("无效的步态类型，使用默认的三角步态")
    
    def set_direction(self, direction):
        """设置移动方向：'forward', 'backward', 'left', 'right'"""
        if direction in ["forward", "backward", "left", "right"]:
            self.direction = direction
        else:
            print("无效的方向，使用默认的前进方向")
    
    def set_speed(self, speed):
        """设置移动速度 (0-100)"""
        self.speed = max(0, min(100, speed))
    
    def stand(self):
        """让机器人站立"""
        for leg in range(6):
            base_positions = self.center_positions[leg*3:(leg+1)*3].copy()
            # 调整站立高度
            base_positions[1] += self.step_height * 0.5
            self.set_leg_position(leg, base_positions)
    
    def rest(self):
        """让机器人休息(收起腿)"""
        for leg in range(6):
            base_positions = self.center_positions[leg*3:(leg+1)*3].copy()
            # 调整休息位置
            base_positions[1] -= self.step_height * 0.3
            self.set_leg_position(leg, base_positions)
    
    def walk(self, steps=1, blocking=True):
        """执行行走动作
        steps: 行走步数
        blocking: 是否阻塞直到完成
        """
        # 计算每步的相位增量
        phase_increment = 0.02
        steps_per_cycle = int(1 / phase_increment)
        
        # 计算总步数
        total_steps = steps * steps_per_cycle
        
        # 执行步态循环
        for i in range(total_steps):
            phase = (i * phase_increment) % 1
            
            # 根据选定的步态类型执行相应的步态函数
            if self.gait_type == "tripod":
                self.tripod_gait(phase)
            else:
                self.wave_gait(phase)
            
            # 根据速度调整延迟
            delay = 10 - (self.speed / 10)
            delay = max(1, delay)  # 确保最小延迟为1ms
            time.sleep_ms(int(delay))
        
        # 回到站立位置
        self.stand()